package com.cy.ragbase.service;

import com.cy.ragbase.dto.*;
import com.cy.ragbase.entity.CodeBlock;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.EmbeddingModel;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.time.LocalDateTime;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

@Service
@RequiredArgsConstructor
@Slf4j
public class CodeBlockService {
    
    private final EmbeddingModel embeddingModel;

    private final MilvusService milvusService;
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final RedisTemplate redisTemplate;



    // 临时存储，生产环境应该使用数据库
    private final Map<String, CodeBlock> codeBlockStorage = new HashMap<>();
    
    /**
     * 批量导入代码块
     */
    public CodeBlockImportResponse importCodeBlocks(CodeBlockImportRequest request) {
        long startTime = System.currentTimeMillis();
        
        List<String> importedIds = new ArrayList<>();
        List<String> errors = new ArrayList<>();
        int successCount = 0;
        int failureCount = 0;
        
        try {
            for (CodeBlockImportRequest.CodeBlockData blockData : request.getCodeBlocks()) {
                try {
                    CodeBlock codeBlock = convertToCodeBlock(blockData);
                    
                    // 生成向量表示
                    String textForEmbedding = buildTextForEmbedding(codeBlock);
                    Embedding embedding = embeddingModel.embed(textForEmbedding).content();
                    codeBlock.setVector(embedding.vector());
                    
                    // 存储到内存（生产环境应存储到数据库）
                    codeBlockStorage.put(codeBlock.getId(), codeBlock);
                    
                    // 存储向量到Milvus（这里需要扩展MilvusService来支持代码块）
                    // milvusService.insertCodeBlock(codeBlock);
                    
                    importedIds.add(codeBlock.getId());
                    successCount++;
                    
                    log.info("Successfully imported code block: {} - {}", codeBlock.getId(), codeBlock.getName());
                    
                } catch (Exception e) {
                    failureCount++;
                    errors.add("Failed to import code block '" + blockData.getName() + "': " + e.getMessage());
                    log.error("Failed to import code block: {}", blockData.getName(), e);
                }
            }
            
            long processingTime = System.currentTimeMillis() - startTime;
            
            CodeBlockImportResponse response = new CodeBlockImportResponse();
            response.setSuccessCount(successCount);
            response.setFailureCount(failureCount);
            response.setTotalCount(request.getCodeBlocks().size());
            response.setProcessingTimeMs(processingTime);
            response.setImportedIds(importedIds);
            response.setErrors(errors);
            response.setStatus(failureCount == 0 ? "SUCCESS" : "PARTIAL_SUCCESS");
            response.setMessage(String.format("导入完成：成功 %d 个，失败 %d 个", successCount, failureCount));
            
            log.info("Code block import completed: success={}, failure={}, time={}ms", 
                    successCount, failureCount, processingTime);
            
            return response;
            
        } catch (Exception e) {
            log.error("Code block import failed", e);
            throw new RuntimeException("代码块导入失败", e);
        }
    }
    
    /**
     * 搜索代码块
     */

    public CodeSearchResponse searchCodeBlocks(CodeSearchRequest request) {
        long startTime = System.currentTimeMillis();

        try {
            // 生成查询向量
            Embedding queryEmbedding = embeddingModel.embed(request.getQuery()).content();
            float[] queryVector = queryEmbedding.vector();

            // 直接在Milvus中进行高效向量搜索
            List<CodeSearchResponse.CodeSearchResult> results = milvusService.searchSimilarCodeBlocks(
                    queryVector,
                    request.getLanguage(),
                    request.getLabels(),
                    request.getTopK(),
                    request.getThreshold()
            );

            long processingTime = System.currentTimeMillis() - startTime;

            CodeSearchResponse response = new CodeSearchResponse();
            response.setQuery(request.getQuery());
            response.setResults(results);
            response.setTotalResults(results.size());
            response.setProcessingTimeMs(processingTime);

            log.info("代码块搜索完成: query='{}', results={}, time={}ms",
                    request.getQuery(), results.size(), processingTime);
            return response;

        } catch (Exception e) {
            log.error("代码块搜索失败: {}", request.getQuery(), e);
            throw new RuntimeException("代码块搜索失败", e);
        }
    }



    /**
     * 生成代码RAG提示词
     */
    public CodeRagResponse generateCodeRagPrompt(CodeRagRequest request) {
        long startTime = System.currentTimeMillis();

        try {
            // 搜索相关代码块
            CodeSearchRequest searchRequest = new CodeSearchRequest();
            searchRequest.setQuery(request.getRequirement());
            searchRequest.setLanguage(request.getTargetLanguage());
            searchRequest.setLabels(request.getHints());
            searchRequest.setTopK(request.getTopK());
            searchRequest.setThreshold(request.getThreshold());

            CodeSearchResponse searchResponse = searchCodeBlocks(searchRequest);

            // 转换为参考代码并去重
            List<CodeRagResponse.ReferenceCode> referenceCodes = searchResponse.getResults()
                    .stream()
                    .map(this::convertToReferenceCode)
                    .collect(Collectors.toList());
            
            // 去重处理：基于代码内容去重，保留相似度最高的
            referenceCodes = deduplicateReferenceCodes(referenceCodes);

            // 生成提示词
            String generatedPrompt = buildCodeGenerationPrompt(request, referenceCodes);

            long processingTime = System.currentTimeMillis() - startTime;

            CodeRagResponse response = new CodeRagResponse();
            response.setOriginalRequirement(request.getRequirement());
            response.setGeneratedPrompt(generatedPrompt);
            response.setReferenceCodes(referenceCodes);
            response.setTotalRetrieved(referenceCodes.size());
            response.setProcessingTimeMs(processingTime);

            log.info("代码RAG提示词生成完成: requirement='{}', references={}, time={}ms",
                    request.getRequirement(), referenceCodes.size(), processingTime);

            return response;

        } catch (Exception e) {
            log.error("代码RAG提示词生成失败: {}", request.getRequirement(), e);
            throw new RuntimeException("代码RAG提示词生成失败", e);
        }
    }


    private CodeRagResponse.ReferenceCode convertToReferenceCode(CodeSearchResponse.CodeSearchResult result) {
        CodeRagResponse.ReferenceCode refCode = new CodeRagResponse.ReferenceCode();
        refCode.setId(result.getId());
        refCode.setName(result.getName());
        refCode.setCode(result.getCode());
        refCode.setInfo(result.getInfo());
        refCode.setLabels(result.getLabels());
        refCode.setLanguage(result.getLanguage());
        refCode.setScore(result.getScore());
        return refCode;
    }

    /**
     * 去重参考代码，基于代码内容去重，保留相似度最高的
     */
    private List<CodeRagResponse.ReferenceCode> deduplicateReferenceCodes(List<CodeRagResponse.ReferenceCode> referenceCodes) {
        Map<String, CodeRagResponse.ReferenceCode> codeMap = new LinkedHashMap<>();
        
        for (CodeRagResponse.ReferenceCode refCode : referenceCodes) {
            String codeContent = normalizeCode(refCode.getCode());
            
            // 如果已存在相同代码内容，比较相似度，保留更高的
            if (codeMap.containsKey(codeContent)) {
                CodeRagResponse.ReferenceCode existing = codeMap.get(codeContent);
                if (refCode.getScore() > existing.getScore()) {
                    codeMap.put(codeContent, refCode);
                }
            } else {
                codeMap.put(codeContent, refCode);
            }
        }
        
        return new ArrayList<>(codeMap.values());
    }

    /**
     * 标准化代码内容，用于去重比较
     */
    private String normalizeCode(String code) {
        if (code == null) return "";
        
        // 去除多余空白字符和换行符，保留基本结构
        return code.replaceAll("\\s+", " ").trim();
    }

    /**
     * 将请求数据转换为CodeBlock实体
     */
    private CodeBlock convertToCodeBlock(CodeBlockImportRequest.CodeBlockData blockData) {
        return CodeBlock.builder()
                .id(UUID.randomUUID().toString())
                .name(blockData.getName())
                .code(blockData.getCode())
                .info(blockData.getInfo())
                .labels(blockData.getLabel())
                .language(blockData.getLanguage())
                .createTime(LocalDateTime.now())
                .updateTime(LocalDateTime.now())
                .build();
    }
    
    /**
     * 构建用于向量化的文本
     */
    private String buildTextForEmbedding(CodeBlock codeBlock) {
        StringBuilder text = new StringBuilder();
        text.append("模块名称: ").append(codeBlock.getName()).append("\n");
        text.append("描述: ").append(codeBlock.getInfo()).append("\n");
        text.append("标签: ").append(String.join(", ", codeBlock.getLabels())).append("\n");
        text.append("编程语言: ").append(codeBlock.getLanguage()).append("\n");
        text.append("代码内容: ").append(codeBlock.getCode());
        return text.toString();
    }
    
    /**
     * 检查代码块是否匹配过滤条件
     */
    private boolean matchesFilters(CodeBlock codeBlock, CodeSearchRequest request) {
        // 语言过滤
        if (request.getLanguage() != null && !request.getLanguage().isEmpty()) {
            if (!request.getLanguage().equalsIgnoreCase(codeBlock.getLanguage())) {
                return false;
            }
        }
        
        // 标签过滤
        if (request.getLabels() != null && !request.getLabels().isEmpty()) {
            boolean hasMatchingLabel = false;
            for (String requestLabel : request.getLabels()) {
                for (String codeLabel : codeBlock.getLabels()) {
                    if (codeLabel.toLowerCase().contains(requestLabel.toLowerCase())) {
                        hasMatchingLabel = true;
                        break;
                    }
                }
                if (hasMatchingLabel) break;
            }
            if (!hasMatchingLabel) {
                return false;
            }
        }
        
        return true;
    }
    
    /**
     * 计算余弦相似度
     */
    private float calculateCosineSimilarity(float[] vectorA, float[] vectorB) {
        if (vectorA.length != vectorB.length) {
            throw new IllegalArgumentException("Vector dimensions must match");
        }
        
        float dotProduct = 0.0f;
        float normA = 0.0f;
        float normB = 0.0f;
        
        for (int i = 0; i < vectorA.length; i++) {
            dotProduct += vectorA[i] * vectorB[i];
            normA += vectorA[i] * vectorA[i];
            normB += vectorB[i] * vectorB[i];
        }
        
        return dotProduct / (float) (Math.sqrt(normA) * Math.sqrt(normB));
    }
    
    /**
     * 构建代码生成提示词
     */
    private String buildCodeGenerationPrompt(CodeRagRequest request, List<CodeRagResponse.ReferenceCode> referenceCodes) {
        StringBuilder prompt = new StringBuilder();

        prompt.append("# 代码生成任务\n\n");
        prompt.append("## 开发需求\n");
        prompt.append(request.getRequirement()).append("\n\n");

        if (request.getTargetLanguage() != null && !request.getTargetLanguage().isEmpty()) {
            prompt.append("## 目标编程语言\n");
            prompt.append(request.getTargetLanguage()).append("\n\n");
        }

        if (request.getHints() != null && !request.getHints().isEmpty()) {
            prompt.append("## 相关提示\n");
            prompt.append("- ").append(String.join("\n- ", request.getHints())).append("\n\n");
        }

        if (!referenceCodes.isEmpty()) {
            prompt.append("## 参考代码示例\n\n");

            for (int i = 0; i < referenceCodes.size(); i++) {
                CodeRagResponse.ReferenceCode refCode = referenceCodes.get(i);
                prompt.append("### 参考示例 ").append(i + 1)
                        .append(": ").append(refCode.getName())
                        .append(" (相似度: ").append(String.format("%.2f", refCode.getScore())).append(")\n\n");

                prompt.append("**功能描述:** ").append(refCode.getInfo()).append("\n\n");

                if (refCode.getLabels() != null && !refCode.getLabels().isEmpty()) {
                    prompt.append("**标签:** ").append(String.join(", ", refCode.getLabels())).append("\n\n");
                }

                prompt.append("**语言:** ").append(refCode.getLanguage()).append("\n\n");

                prompt.append("**代码:**\n");
                prompt.append("```").append(refCode.getLanguage().toLowerCase()).append("\n");
                prompt.append(refCode.getCode());
                prompt.append("\n```\n\n");
                prompt.append("---\n\n");
            }
        }

        prompt.append("## 代码生成要求\n");

        switch (request.getCodeStyle()) {
            case PRODUCTION:
                prompt.append("请生成生产级别的代码，要求：\n");
                prompt.append("1. 代码健壮性强，包含完整的错误处理\n");
                prompt.append("2. 遵循最佳实践和设计模式\n");
                prompt.append("3. 代码可读性高，结构清晰\n");
                prompt.append("4. 包含必要的性能优化\n");
                break;

            case PROTOTYPE:
                prompt.append("请生成原型级别的代码，要求：\n");
                prompt.append("1. 快速实现核心功能\n");
                prompt.append("2. 代码简洁明了\n");
                prompt.append("3. 专注于主要逻辑实现\n");
                break;

            case EDUCATIONAL:
                prompt.append("请生成教学级别的代码，要求：\n");
                prompt.append("1. 详细的注释说明每个步骤\n");
                prompt.append("2. 清晰的变量和函数命名\n");
                prompt.append("3. 包含原理解释和使用示例\n");
                prompt.append("4. 易于理解和学习\n");
                break;
        }

        prompt.append("5. 参考上述示例代码的实现思路和风格\n");
        prompt.append("6. 确保代码的正确性和完整性\n\n");

        prompt.append("请根据需求和参考示例，生成符合要求的代码实现。");

        return prompt.toString();
    }


    /**
     * 从JSON文件批量导入代码块
     */
    public CodeBlockImportResponse importCodeBlocksFromFile(MultipartFile file) {
        long startTime = System.currentTimeMillis();

        List<String> importedIds = new ArrayList<>();
        List<String> errors = new ArrayList<>();
        int successCount = 0;
        int failureCount = 0;
        int totalCount = 0;

        try {
            // 读取并解析JSON文件
            String jsonContent = new String(file.getBytes(), "UTF-8");
            List<CodeBlockImportRequest.CodeBlockData> codeBlockDataList = parseJsonFile(jsonContent);

            totalCount = codeBlockDataList.size();
            log.info("开始导入 {} 个代码块，文件: {}", totalCount, file.getOriginalFilename());

            List<CodeBlock> codeBlocksToInsert = new ArrayList<>();

            for (CodeBlockImportRequest.CodeBlockData blockData : codeBlockDataList) {
                try {
                    if (!isValidCodeBlockData(blockData)) {
                        failureCount++;
                        errors.add("代码块数据不完整: " + blockData.getName());
                        continue;
                    }

                    CodeBlock codeBlock = convertToCodeBlock(blockData);

                    // 生成向量表示
                    String textForEmbedding = buildTextForEmbedding(codeBlock);
                    Embedding embedding = embeddingModel.embed(textForEmbedding).content();
                    codeBlock.setVector(embedding.vector());

                    // 保存元数据到Redis（可选，用于快速查询元信息）
                    saveCodeBlockMetadataToRedis(codeBlock);

                    codeBlocksToInsert.add(codeBlock);
                    importedIds.add(codeBlock.getId());
                    successCount++;

                    log.debug("准备导入代码块: {} - {}", codeBlock.getId(), codeBlock.getName());

                } catch (Exception e) {
                    failureCount++;
                    String errorMsg = "导入代码块失败 '" +
                            (blockData.getName() != null ? blockData.getName() : "Unknown") +
                            "': " + e.getMessage();
                    errors.add(errorMsg);
                    log.error("导入代码块失败: {}", blockData.getName(), e);
                }
            }

            // 批量插入到Milvus（核心存储）- 分批处理避免gRPC消息过大
            if (!codeBlocksToInsert.isEmpty()) {
                insertCodeBlocksInBatches(codeBlocksToInsert);
                log.info("成功向Milvus插入 {} 个代码块", codeBlocksToInsert.size());
            }

            long processingTime = System.currentTimeMillis() - startTime;

            CodeBlockImportResponse response = new CodeBlockImportResponse();
            response.setSuccessCount(successCount);
            response.setFailureCount(failureCount);
            response.setTotalCount(totalCount);
            response.setProcessingTimeMs(processingTime);
            response.setImportedIds(importedIds);
            response.setErrors(errors);

            if (failureCount == 0) {
                response.setStatus("SUCCESS");
                response.setMessage(String.format("成功导入 %d 个代码块", successCount));
            } else if (successCount > 0) {
                response.setStatus("PARTIAL_SUCCESS");
                response.setMessage(String.format("部分成功：成功 %d 个，失败 %d 个", successCount, failureCount));
            } else {
                response.setStatus("FAILED");
                response.setMessage(String.format("导入失败：%d 个代码块全部失败", totalCount));
            }

            log.info("代码块导入完成: 文件={}, 成功={}, 失败={}, 时间={}ms",
                    file.getOriginalFilename(), successCount, failureCount, processingTime);

            return response;

        } catch (Exception e) {
            log.error("代码块导入失败", e);
            throw new RuntimeException("代码块导入失败: " + e.getMessage(), e);
        }
    }


    private void saveCodeBlockMetadataToRedis(CodeBlock codeBlock) {
        try {
            String key = "codeblock_meta:" + codeBlock.getId();
            redisTemplate.opsForHash().put(key, "id", codeBlock.getId());
            redisTemplate.opsForHash().put(key, "name", codeBlock.getName());
            redisTemplate.opsForHash().put(key, "language", codeBlock.getLanguage());
            redisTemplate.opsForHash().put(key, "labels", String.join(",", codeBlock.getLabels()));
            redisTemplate.opsForHash().put(key, "createTime", codeBlock.getCreateTime().toString());

            redisTemplate.expire(key, 30, TimeUnit.DAYS);
        } catch (Exception e) {
            log.warn("保存代码块元数据到Redis失败: {}", codeBlock.getId(), e);
            // 这不是致命错误，不影响主要流程
        }
    }

    /**
     * 解析JSON文件内容
     */
    private List<CodeBlockImportRequest.CodeBlockData> parseJsonFile(String jsonContent) {
        try {
            // 直接解析为CodeBlockData数组
            CodeBlockImportRequest.CodeBlockData[] dataArray =
                    objectMapper.readValue(jsonContent, CodeBlockImportRequest.CodeBlockData[].class);
            return Arrays.asList(dataArray);

        } catch (Exception e) {
            log.error("Failed to parse JSON content", e);
            throw new RuntimeException("JSON格式解析失败，请检查文件格式: " + e.getMessage(), e);
        }
    }



    /**
     * 验证代码块数据是否有效
     */
    private boolean isValidCodeBlockData(CodeBlockImportRequest.CodeBlockData blockData) {
        return blockData != null &&
                blockData.getName() != null && !blockData.getName().trim().isEmpty() &&
                blockData.getCode() != null && !blockData.getCode().trim().isEmpty() &&
                blockData.getLanguage() != null && !blockData.getLanguage().trim().isEmpty();
    }

    /**
     * 获取所有已导入的代码块（用于调试和管理）
     */
    public List<CodeBlock> getAllCodeBlocks() {
        return new ArrayList<>(codeBlockStorage.values());
    }

    /**
     * 根据ID获取代码块
     */
    public CodeBlock getCodeBlockById(String id) {
        return codeBlockStorage.get(id);
    }

    /**
     * 删除代码块
     */
    public boolean deleteCodeBlock(String id) {
        return codeBlockStorage.remove(id) != null;
    }

    /**
     * 获取代码块统计信息
     */
    public Map<String, Object> getCodeBlockStats() {
        Map<String, Object> stats = new HashMap<>();
        stats.put("totalCount", codeBlockStorage.size());

        // 按语言统计
        Map<String, Long> languageStats = codeBlockStorage.values().stream()
                .collect(Collectors.groupingBy(CodeBlock::getLanguage, Collectors.counting()));
        stats.put("byLanguage", languageStats);

        // 按标签统计
        Map<String, Long> labelStats = new HashMap<>();
        for (CodeBlock codeBlock : codeBlockStorage.values()) {
            if (codeBlock.getLabels() != null) {
                for (String label : codeBlock.getLabels()) {
                    labelStats.merge(label, 1L, Long::sum);
                }
            }
        }
        stats.put("byLabel", labelStats);

        return stats;
    }

    /**
     * 分批插入代码块到Milvus，避免gRPC消息过大
     */
    private void insertCodeBlocksInBatches(List<CodeBlock> codeBlocks) {
        final int BATCH_SIZE = 100; // 每批处理100个代码块
        
        int totalBatches = (int) Math.ceil((double) codeBlocks.size() / BATCH_SIZE);
        log.info("开始分批插入 {} 个代码块，分为 {} 批，每批 {} 个", 
                codeBlocks.size(), totalBatches, BATCH_SIZE);
        
        for (int i = 0; i < codeBlocks.size(); i += BATCH_SIZE) {
            int endIndex = Math.min(i + BATCH_SIZE, codeBlocks.size());
            List<CodeBlock> batch = codeBlocks.subList(i, endIndex);
            
            int currentBatch = (i / BATCH_SIZE) + 1;
            
            try {
                milvusService.insertCodeBlocks(batch);
                log.info("批次 {}/{} 插入成功: {} 个代码块", 
                        currentBatch, totalBatches, batch.size());
                
                // 短暂休息，避免过于频繁的请求
                Thread.sleep(50);
                
            } catch (Exception e) {
                log.error("批次 {}/{} 插入失败: {}", currentBatch, totalBatches, e.getMessage(), e);
                throw new RuntimeException("批量插入失败在第 " + currentBatch + " 批: " + e.getMessage(), e);
            }
        }
        
        log.info("所有批次插入完成，共插入 {} 个代码块", codeBlocks.size());
    }

    private void saveCodeBlockToRedis(CodeBlock codeBlock) {
        try {
            String key = "codeblock:" + codeBlock.getId();
            redisTemplate.opsForHash().put(key, "id", codeBlock.getId());
            redisTemplate.opsForHash().put(key, "name", codeBlock.getName());
            redisTemplate.opsForHash().put(key, "language", codeBlock.getLanguage());
            redisTemplate.opsForHash().put(key, "labels", String.join(",", codeBlock.getLabels()));
            redisTemplate.opsForHash().put(key, "createTime", codeBlock.getCreateTime().toString());

            redisTemplate.expire(key, 30, TimeUnit.DAYS);

        } catch (Exception e) {
            log.error("保存代码块到Redis失败: {}", codeBlock.getId(), e);
        }
    }
} 